#!/usr/bin/env python
# -*- coding:utf-8 -*-
# auth: wuye

import hmac
import hashlib
import time
import struct
import base64


def base32_decode(encoded):
    """Base32解码，确保密钥长度是8的倍数(添加填充)"""
    encoded = encoded.upper()
    padding = 8 - (len(encoded) % 8)
    if padding != 8:
        encoded += '=' * padding
    try:
        return base64.b32decode(encoded)
    except Exception as e:
        print(f"Base32解码错误: {e}")
        return None


def generate_totp(secret, timestamp=None, interval=30, digits=6) -> str:
    """
    生成基于时间的一次性密码（TOTP）
    参数:
        secret: 共享密钥（未经过Base32编码）
        timestamp: 时间戳（如果为None，则使用当前时间）
        interval: 时间间隔（秒）
        digits: TOTP代码的位数

    返回值:
        str: 生成的TOTP代码
    """
    # 使用Base32解码密钥
    key = base32_decode(secret)
    print(f"十六进制key: {key.hex()}")

    # 计算时间步长
    if timestamp is None:
        timestamp = int(time.time())
    timestep = int(timestamp // interval)

    # 将时间步长转换为8字节大端字节串
    # 格式字符串">Q"中，>表示大端序（网络字节序），Q：无符号长整数（8字节），I：无符号整数（4字节）。
    # 之所以使用8字节(64位)，是因为TOTP标准规定时间步长应为64位整数。这足够覆盖当前Unix时间戳所有可预见的时间范围。
    msg = struct.pack(">Q", timestep)

    # 计算HMAC-SHA1
    hmac_hash = hmac.new(key, msg, hashlib.sha1).digest()

    # 动态截断
    offset = hmac_hash[-1] & 0x0F
    truncated_hash = hmac_hash[offset:offset + 4]

    # 将截断的字节转换为整数
    code_int = struct.unpack(">I", truncated_hash)[0] & 0x7FFFFFFF

    # 计算TOTP值
    totp_value = code_int % (10 ** digits)

    # 添加前导零使长度固定
    return str(totp_value).zfill(digits)


def verify_totp(secret, code, interval=30, digits=6) -> bool:
    """
    验证TOTP代码
    参数:
        secret: 共享密钥（未经过Base32编码）
        code: 要验证的TOTP代码
        digits: TOTP代码的位数

    返回值:
        bool: 如果代码有效则返回True，否则返回False
    """
    # 考虑到时间漂移，检查前后一个时间窗口(总共3个窗口)
    timestamp = int(time.time())

    # 当前时间窗口
    if generate_totp(secret, timestamp, digits=digits) == code:
        return True
    # 前一个时间窗口（x秒之前）
    if generate_totp(secret, timestamp - interval, digits=digits) == code:
        return True
    # 后一个时间窗口（x秒之后）
    if generate_totp(secret, timestamp + interval, digits=digits) == code:
        return True

    return False


def generate_totp2(secret=None):
    """使用三方库更简单的实现"""
    import pyotp

    # 生成共享密钥
    if secret is None:
        secret = pyotp.random_base32()
    print(f"共享密钥: {secret}")

    # 创建TOTP对象
    totp = pyotp.TOTP(secret, interval=30, digits=6)

    # 生成当前验证码
    current_code = totp.now()
    print(f"当前TOTP验证码: {current_code}")

    # 验证验证码
    print(f"验证结果: {totp.verify(current_code)}")

    # 生成Provisioning URI (用于生成二维码)
    uri = totp.provisioning_uri("user@example.com", issuer_name="MyApp")
    print(f"\nProvisioning URI:\n{uri}")


# Example usage
if __name__ == "__main__":
    test_secret = "IU655SJR7ISY"
    ret = generate_totp(test_secret)
    print(f"genertae_totp生成的TOTP为：{ret}")
    generate_totp2(test_secret)
